package cn.i7baozh.starter.common.interceptor;

import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

/**
 * @author : baoqi.zhang
 * @date :   2022/5/12 14:10
 */
@Slf4j
public class RequestLoggingInterceptor implements HandlerInterceptor, WebMvcConfigurer {

    private static final String DISTRIBUTE_LOGGING_FLAG = "_traceId_";

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String header = request.getHeader(DISTRIBUTE_LOGGING_FLAG);
        if (header != null && header.length() > 0) {
            MDC.put(DISTRIBUTE_LOGGING_FLAG, header);
        }
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        MDC.remove(DISTRIBUTE_LOGGING_FLAG);
    }

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(this);
    }
}
